import os

from openai import OpenAI

# OneAPI的令牌
ONE_API_API_KEY = os.getenv("ONE_API_API_KEY")
# OneAPI服务的地址
ONE_API_API_BASE_URL = "http://localhost:3000/v1"


def call_llm(model: str):
    """
    调用大语言模型
    :param model:
    :return:
    """

    client = OpenAI(api_key=ONE_API_API_KEY, base_url=ONE_API_API_BASE_URL)

    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": "You are a helpful assistant"},
            {"role": "user", "content": "杭州是个风水宝地吗？"},
        ],
        stream=False
    )

    print(response.choices[0].message.content)


if __name__ == '__main__':
    # DeepSeek
    # model = "deepseek-chat"
    # 通义
    model = "qwen-plus"
    call_llm(model)
